/*
* Copyright (c) Microsoft Corporation.  All Rights Reserved.
*
* Public Mock10 API
*/
#pragma once
#include "Mock10.h"
#include "MockFunction.h"
#include "MockStatus.h"
#include "WexTypes.h"

#pragma push_macro("Assert")
#undef Assert

#pragma push_macro("Trace")
#undef Trace

#pragma push_macro("SEALED")
#undef SEALED

#if _MSC_VER >= 1400
# define SEALED sealed
#else
# define SEALED
#endif

namespace Mock10
{
    class DynamicVariable;

    template <typename TVariableType>
    class MockVariable;

    namespace Private
    {
        class MOCK10_API EnableWrite
        {
        public:
            EnableWrite(void* pAddress, size_t size);
            ~EnableWrite();

            HRESULT GetStatus();

        private:
            void* m_pAddress;
            unsigned long m_previous;
            size_t m_size;
            HRESULT m_status;
        };

        class MOCK10_API MockVariableBase
        {
        friend class DynamicVariableList;

        template <typename T>
        friend class CachedVariable;

        template <typename TVariableType>
        friend class MockVariable;

        public:
            MockVariableBase& operator= (const MockVariableBase& other);

            HRESULT GetStatus() const
            {
                return m_status;
            }

            size_t GetSize() const;

        protected:
            MockVariableBase(const wchar_t* pszTargetBinary, const wchar_t* pszTargetQuery);
            MockVariableBase(const MockVariableBase& other);
            MockVariableBase();

            // wchar_t native type exports
            #if defined(MOCK10_FULL_BUILD)
            MockVariableBase(const __wchar_t* pszTargetBinary, const __wchar_t* pszTargetQuery);
            #endif

            virtual ~MockVariableBase();
            void ValidateType(const void* pVariable);
            void* GetAddress() const;

            HRESULT m_status;

        private:
            static void __stdcall Trace(const wchar_t* pszMessage);
            // wchar_t native type exports
            #if defined(MOCK10_FULL_BUILD)
            static void __stdcall Trace(const __wchar_t* pszMessage);
            #endif

            const wchar_t* GetName() const;
            void Initialize(const wchar_t* pszTargetBinary, const wchar_t* pszTargetQuery);

            struct Impl;
            Impl* m_pImpl;
        };

        template <typename T>
        struct IsBlittable
        {
            static const bool Result = __is_pod(T) || !__is_class(T);
        };

        template<typename T, unsigned int N>
        struct IsBlittable<T[N]>
        {
            static const bool Result = __is_pod(T) || !__is_class(T);
        };

        template <typename T, bool IsBlittable>
        struct SetDataImpl;

        // Blittable types (arrays or single values)
        template<typename T>
        struct SetDataImpl<T, true>
        {
            static void Work(const T& value, T& destination)
            {
                WEX_ASSERT(sizeof(value) == sizeof(destination), L"");
                ::memcpy_s(&destination, sizeof(T), &value, sizeof(T));
            }
        };

        // Non-blittable single value
        template <typename T>
        struct SetDataImpl<T, false>
        {
            static void Work(const T& value, T& destination)
            {
                destination = value;
            }
        };

        // Non-blittable array
        template <typename T, size_t Count>
        struct SetDataImpl<T[Count], false>
        {
            static void Work(const T(&value)[Count], T(&destination)[Count])
            {
                for (size_t i = 0; i < Count; ++i)
                {
                    destination[i] = value[i];
                }
            }
        };

        template<typename T>
        void SetData(const T& value, T& destination)
        {
            SetDataImpl<T, IsBlittable<T>::Result>::Work(value, destination);
        }

        template <typename T>
        class CachedVariable
        {
        public:
            CachedVariable()
                : m_pAddress(NULL)
            {
            }

            void Save(T& value)
            {
                SetData(value, m_storage);
                m_pAddress = &value; // Store m_pAddress after the call to Store, in case the call throws
            }

            bool IsValid() const
            {
                return m_pAddress != NULL;
            }

            const T& GetValue() const
            {
                return m_storage;
            }

            ~CachedVariable()
            {
                if (m_pAddress)
                {
                    EnableWrite enable(m_pAddress, sizeof(T));
                    WEX_ASSERT(SUCCEEDED(enable.GetStatus()), L"~DynamicVariableImpl is about to write to potentially un-writable memory");

                    #if !defined(NO_MOCK_THROW)
                    try
                    {
                    #endif
                        SetData(m_storage, *reinterpret_cast<T*>(m_pAddress));
                    #if !defined(NO_MOCK_THROW)
                    }
                    catch (...)
                    {
                        MockVariableBase::Trace(L"Mock10: Failed to restore the mocked variable value(s).");
                    }
                    #endif
                }
            }

        private:
            CachedVariable& operator= (const CachedVariable& other); // Not implemented

            void* m_pAddress;
            T m_storage;
        };

        template <typename T>
        struct __declspec(dllexport) PinSymbol
        {
            static const T* m_pT;
        };

        template <typename T>
        const T* PinSymbol<T>::m_pT = NULL;

        class MOCK10_API DynamicVariableList
        {
        friend class Mock10::DynamicVariable;

        public:
            DynamicVariableList& operator= (const DynamicVariableList& other);

        private:
            DynamicVariableList();
            DynamicVariableList(const DynamicVariableList& other);
            ~DynamicVariableList();

            DynamicVariable& Get(const wchar_t* pszName); 

            // wchar_t native type exports
            #if defined(MOCK10_FULL_BUILD)
            DynamicVariable& Get(const __wchar_t* pszName);
            #endif

            void Initialize(DynamicVariable& dynamicVariable);

            HRESULT m_status;
            struct Impl;
            Impl* m_pImpl;
        };

        MOCK10_API const DynamicVariable& GetEmptyDynamicVariable();
    }

    template <typename TVariableType>
    class MockVariable: Private::MockVariableBase
    {
    friend class DynamicVariable;

    public:
        MockVariable(const wchar_t* pszTargetBinary, const wchar_t* pszTargetQuery)
            : Private::MockVariableBase(pszTargetBinary, pszTargetQuery)
        {
            if (SUCCEEDED(m_status))
            {
                ValidateType(reinterpret_cast<const void*>(&Private::PinSymbol<TVariableType>::m_pT));
            }

            #if !defined(NO_MOCK_THROW)
            if (FAILED(m_status))
            {
                Private::Throw::MockException(m_status);
            }
            #endif
        }

        MockVariable(const MockVariable& other)
            : MockVariableBase(other)
        {
            if (SUCCEEDED(m_status))
            {
                m_originalValue.Save(*reinterpret_cast<TVariableType*>(GetAddress()));
            }
        }

        void SetValue(const TVariableType& value)
        {
            if (SUCCEEDED(m_status))
            {
                Private::EnableWrite enable(GetAddress(), GetSize());
                m_status = enable.GetStatus();
                if (SUCCEEDED(m_status))
                {
                    if (!m_originalValue.IsValid())
                    {
                    #if !defined(NO_MOCK_THROW)
                        try
                        {
                    #endif
                            m_originalValue.Save(*reinterpret_cast<TVariableType*>(GetAddress()));
                    #if !defined(NO_MOCK_THROW)
                        }
                        catch (...)
                        {
                            Private::MockVariableBase::Trace(L"Mock10: Failed to store the mocked variable value(s).");
                            m_status = E_UNEXPECTED;
                            throw;
                        }
                    #endif
                    }

                #if !defined(NO_MOCK_THROW)
                    try
                    {
                #endif
                        Private::SetData<TVariableType>(value, *reinterpret_cast<TVariableType*>(GetAddress()));
                #if !defined(NO_MOCK_THROW)
                    }
                    catch (...)
                    {
                        Private::MockVariableBase::Trace(L"Mock10: Failed to set the mocked variable value(s).");
                        m_status = E_UNEXPECTED;
                        throw;
                    }
                #endif
                }
            }

        #if !defined(NO_MOCK_THROW)
            if (FAILED(m_status))
            {
                Private::Throw::MockException(m_status);
            }
        #endif
        }
            
        const TVariableType& GetValue() const
        {
            if (SUCCEEDED(m_status))
            {
                return *reinterpret_cast<TVariableType*>(GetAddress());
            }

            #if !defined(NO_MOCK_THROW)
                if (FAILED(m_status))
                {
                    Private::Throw::MockException(m_status);
                }
            #endif

            // For non-throwing clients, we return a reference to the stored value (or an empty object if nothing was stored yet)
            // The user needs to call GetStatus() before making this call
            return m_originalValue.GetValue();
        }

        using Private::MockVariableBase::GetStatus;

    private:
        explicit MockVariable(const DynamicVariable& dynamicVariable)
            : Private::MockVariableBase(dynamicVariable)
        {
            if (SUCCEEDED(m_status))
            {
                ValidateType(reinterpret_cast<const void*>(&Private::PinSymbol<TVariableType>::m_pT));
            }

        #if !defined(NO_MOCK_THROW)
            if (FAILED(m_status))
            {
                Private::Throw::MockException(m_status);
            }
        #endif
        }

        MockVariable& operator= (const MockVariable& other); // Not implemented

        Private::CachedVariable<TVariableType> m_originalValue;
    };

    class DynamicVariable: Private::MockVariableBase
    {
    template <typename TVariableType>
    friend class MockVariable;
    friend class Private::DynamicVariableList;

    public:
        DynamicVariable(const wchar_t* pszTargetBinary, const wchar_t* pszTargetQuery)
            : Private::MockVariableBase(pszTargetBinary, pszTargetQuery)
        {
            #if !defined(NO_MOCK_THROW)
                if (FAILED(m_status))
                {
                    Private::Throw::MockException(m_status);
                }
            #endif

            if (SUCCEEDED(m_status))
            {
                m_members.Initialize(*this);
            }
        }

        template <typename TVariableType>
        MockVariable<TVariableType> GetMember(const wchar_t* pszName) const
        {
            return MockVariable<TVariableType>(GetMember(pszName));
        }

        const DynamicVariable& GetMember(const wchar_t* pszName) const
        {
            // If our list initialization failed, do not continue
            if (FAILED(m_members.m_status))
            {
            #if !defined(NO_MOCK_THROW)
                Private::Throw::MockException(m_members.m_status);
            #else
                return Private::GetEmptyDynamicVariable();
            #endif
            }

            DynamicVariable& variable = m_members.Get(pszName);
            
            // Throw here if constructing the variable failed
            #if !defined(NO_MOCK_THROW)
                if (FAILED(variable.GetStatus()))
                {
                    Private::Throw::MockException(variable.GetStatus());
                }
            #endif

            return variable;
        };

        DynamicVariable()
            : Private::MockVariableBase()
        {
        }

        explicit DynamicVariable(HRESULT status)
            : Private::MockVariableBase()
        {
            m_status = status;
        }

        using Private::MockVariableBase::GetStatus;
        
    private:
        mutable Private::DynamicVariableList m_members;
    };
} /* namespace Mock10 */

#pragma pop_macro("SEALED")
#pragma pop_macro("Trace")
#pragma pop_macro("Assert")
